
import TrainAndTest


import re            # 正则表达式
import collections   # 计数统计
import random
import math
import torch
from torch import nn
from torch.nn import functional as F


def read_txtfile(txtpath):  
    """将txtpath的txt文件的英文文本加载到文本行的列表
    Args:
        txtpath (str): 文本路径
    Returns:
        list: 列表，列表里的元素是txtpath中的行(非字母元素用空格替代)。
    """

    with open(txtpath, 'r') as f:
        lines = f.readlines()
    """
    re.sub('[^A-Za-z]+', ' ', line)：使用正则表达式将所有非字母的字符替换为空格。 正则表达式中的 + 表示匹配前面的元素一次或多次。返回替换后的新字符串。
    .strip()：去除字符串前后的空白字符。返回新字符串。
    .lower()：将所有字母转换为小写。返回新字符串。
    """
    return [re.sub('[^A-Za-z]+', ' ', line).strip().lower() for line in lines]

def tokenize(lines, token='word'): 
    """将文本行拆分为单词或字符词元.
    Args:
        lines (list): 文本行列表,列表元素是一行文本。
        token (str, optional): 'word'or'char',按照单词或字母作为词元基本单位. Defaults to 'word'.
    Returns:
        list: 返回的列表的元素是字符串列表(字符串列表对应一行文本)，字符串列表的元素是词元，每个词元都是一个字符串。
    """
    if token == 'char':
        return [list(line) for line in lines] # list()函数将每行文本拆分为字符列表,这里空格也算一个词元
    else : #默认word
        return [line.split() for line in lines]  #split()方法将每行文本按空格分割成单词列表

def count_corpus(tokens):  
    """统计词元的频率,corpus:语料库.
    Args:
        tokens (list): 1D列表或2D列表，代表整个语料的词元集合
    Returns:
        字典: 键是词元，值是该词元在tokens中出现的次数。
    """
    if len(tokens) == 0 or isinstance(tokens[0], list):
        # 将词元列表展平成一个列表
        tokens = [token for line in tokens for token in line] # 如果tokens是空的，会展开成空列表。
    #   collections.Counter是 Python 标准库中的一个工具类，专门用于计数可哈希对象
    #   它会返回一个类似字典的对象，其中键是词元，值是该词元在整个语料中出现的次数
    return collections.Counter(tokens)

class Vocab: 

    def __init__(self, tokens=None, min_freq=0, reserved_tokens=None):
        """_summary_
        Args:
            tokens (list, optional): 语料库的词元列表. Defaults to None.
            min_freq (int, optional): 词元最小出现频次(低于这个频次的单词忽略). Defaults to 0.
            reserved_tokens (list, optional): 指定需要在词汇表中保留的特殊词元（也称为 “预留词元”）。这些词元通常具有特殊含义，需要强制包含在词汇表中，不受词频限制。Defaults to None.
            例如：
                \<pad>：用于句子长度补齐（padding）
                \<bos>：表示句子开始（beginning of sentence）
                \<eos>：表示句子结束（end of sentence）
                \<mask>：用于掩码语言模型（如 BERT 中的掩码任务）
        """
        if tokens is None:
            tokens = []
        if reserved_tokens is None:
            reserved_tokens = []
        # 按出现频率排序
        counter = count_corpus(tokens)

        # items() 方法会返回一个包含 (词元，频率) 的键值对迭代器.
        # sorted() 是 Python 内置的排序函数.
        #  key=lambda x: x[1] 表示按照元组的第二个元素（即词频）进行排序.
        #      这里 lambda关键字相当于定义了一个匿名函数，函数输入x,输出x[1],x是函数形参。与上下文无关。
        #      key 的值是 lambda关键字定义的函数，sorted 会调用 key 函数来计算元素排序的依据值。
        #  reverse=True 表示降序排列（频率高的词元排在前面）
        self.min_freq = min_freq
        self.token_freqs = sorted(counter.items(), key=lambda x: x[1],
                                   reverse=True)
        # 未知词元的索引为0,然后是保留词元，然后是频率由高到低的词元。
        self.unk = 0  # 未知词元，索引固定是0
        self.idx_to_token = ['<unk>'] + reserved_tokens  # idx_to_token是存储 “索引到词元” 映射关系的列表
        self.token_to_idx = {token: idx for idx, token in enumerate(self.idx_to_token)}  # token_to_idx是词元到索引的映射，用字典实现。
        # enumerate(list) 会对列表进行遍历,每次循环返回两个值:当前元素在列表中的索引(从0开始)和当前元素的值(即词元本身)

        for token, freq in self.token_freqs:
            if freq < self.min_freq:  # 忽略低频词，_token_freqs已经排序过了，后面都是更低频的词，这里可以跳出循环。
                break
            if token not in self.token_to_idx:  # 如果是新的词元，添加到 idx_to_token 和 token_to_idx
                self.idx_to_token.append(token)
                self.token_to_idx[token] = len(self.idx_to_token) - 1


    def set_txtfile(self,txtpath,reserved_tokens=None):
        """重新设置语料库。将txtpath的文本文件的内容添加到词表(原来的词表作废)统计中。 注意，该文件中的低频词也会被忽视。
        Args:
            txtpath (str): 文本路径
            reserved_tokens (list, optional): 指定需要在词汇表中保留的特殊词元（也称为 “预留词元”）。这些词元通常具有特殊含义，需要强制包含在词汇表中，不受词频限制。Defaults to None.
                例如：
                \<pad>：用于句子长度补齐（padding）
                \<bos>：表示句子开始（beginning of sentence）
                \<eos>：表示句子结束（end of sentence）
                \<mask>：用于掩码语言模型（如 BERT 中的掩码任务）

        """
        lines = read_txtfile(txtpath)
        tokens = tokenize(lines)
        if tokens is None:
            tokens = []
        if reserved_tokens is None:
            reserved_tokens = []

        counter = count_corpus(tokens)
        # print(' counts of the : ',counter['the'])
    

        # items() 方法会返回一个包含 (词元，频率) 的键值对迭代器.
        # sorted() 是 Python 内置的排序函数.
        #  key=lambda x: x[1] 表示按照元组的第二个元素（即词频）进行排序.
        #      这里 lambda关键字相当于定义了一个匿名函数，函数输入x,输出x[1],x是函数形参。与上下文无关。
        #      key 的值是 lambda关键字定义的函数，sorted 会调用 key 函数来计算元素排序的依据值。
        #  reverse=True 表示降序排列（频率高的词元排在前面）

        self.token_freqs = sorted(counter.items(), key=lambda x: x[1],
                                   reverse=True)
        # 未知词元的索引为0,然后是保留词元，然后是频率由高到低的词元。
        self.unk = 0  # 未知词元，索引固定是0
        self.idx_to_token = ['<unk>'] + reserved_tokens  # idx_to_token是存储 “索引到词元” 映射关系的列表
        self.token_to_idx = {token: idx for idx, token in enumerate(self.idx_to_token)}  # token_to_idx是词元到索引的映射，用字典实现。
        # enumerate(list) 会对列表进行遍历,每次循环返回两个值:当前元素在列表中的索引(从0开始)和当前元素的值(即词元本身)

        for token, freq in self.token_freqs:
            if freq < self.min_freq:  # 忽略低频词，_token_freqs已经排序过了，后面都是更低频的词，这里可以跳出循环。
                break
            if token not in self.token_to_idx:  # 如果是新的词元，添加到 idx_to_token 和 token_to_idx
                self.idx_to_token.append(token)
                self.token_to_idx[token] = len(self.idx_to_token) - 1

    # 在自定义类中定义 __len__ 和 __getitem__ 这两个方法，是为了让类的实例具有类似序列（如列表、元组）的行为。

    # 例如：  
    #  vocab = Vocab(...) 
    #  print(len(vocab))   # 这里 len(vocab) 相当于调用 vocab.__len__(), 返回词表的大小。
    def __len__(self):
        return len(self.idx_to_token)

    # __getitem__ 方法实现“索引访问” 的功能。  
    # 作用：允许通过 vocab[tokens] 的方式（类似列表索引）来查询词元对应的索引。
    # 支持单个词元和多个词元（列表 / 元组）的查询。 这里返回词元的索引。
    def __getitem__(self, tokens):
        if not isinstance(tokens, (list, tuple)): # 如果 tokens 不是列表或者元组的类型。
            return self.token_to_idx.get(tokens, self.unk) # get：根据键查找值，找不到词元，就返回默认值self.unk
        return [self.__getitem__(token) for token in tokens] 

    def to_tokens(self, indices):
        """根据索引还原成文本。输出的形状与输入一致。
        Args:
            indices (int , list or tuple): 索引，支持单个整形索引，或者索引列表，或者索引元组。
        Returns:
            (str or list): 索引对应的字符串。 
        """
        if not isinstance(indices, (list, tuple)):
            return self.idx_to_token[indices]
        return [self.idx_to_token[index] for index in indices]

def load_corpus(txtpath,max_tokens=-1):  
    """返回txtpath的词元索引列表和词表.
    Args:
        txtpath (str): 语料的txt文件路径
        max_tokens (int, optional): 提取语料库的最大长度。. Defaults to -1(不限制长度)
    Returns:
        (tuple): (corpus, vocab) corpus是预料库列表，元素是词元的索引。vocab 是词元字典，键是词元字符串，值是词元索引。
    """
    lines = read_txtfile(txtpath)
    tokens = tokenize(lines, 'word')
    vocab = Vocab(tokens)
    # 因为时光机器数据集中的每个文本行不一定是一个句子或一个段落，
    # 所以将所有文本行展平到一个列表中
    corpus = [vocab[token] for line in tokens for token in line]
    if max_tokens > 0:
        corpus = corpus[:max_tokens]
    return corpus, vocab  

def seq_data_iter_random(corpus, batch_size, num_steps):  
    """使用随机抽样生成一个小批量子序列。本质就是在列表corpus中，随机截取num_steps大小的子序列。batch_size 个子序列组成一个批次。将输入样本序列后移一位得到输出样本序列。
    Args:
        corpus (list): 词元列表（即语料库）
        batch_size (int): 批量大小，即每个批次的样本个数。
        num_steps (int): 每个样本(即子序列)的序列长度。
    Returns:
        (tuple): (输入张量，输出张量) 张量形状：[batch_size,样本尺寸(num_steps)]
    Yields:
        (tuple): (输入张量，输出张量) 张量形状：[batch_size,样本尺寸(num_steps)]
    """
    # 从随机偏移量开始对序列进行分区，随机范围包括num_steps-1
    # 从原始语料库的前 num_steps 个位置中随机选择一个起始点，对语料库进行截断。避免每次训练都从相同位置开始，增加数据的随机性。
    # 新 corpus 就是原 corpus 前 num_steps 中任意位置开始的子序列。一般num_steps 远小于 corpus 的长度，因此不会对产生的批次数量有影响。
    corpus = corpus[random.randint(0, num_steps - 1):]

    # num_subseqs 表示子序列的数量，每个子序列的长度是 num_steps 
    # 减去1，是因为我们需要考虑标签，然后用地板除。
    num_subseqs = (len(corpus) - 1) // num_steps
    # 长度为num_steps的子序列的起始索引
    initial_indices = list(range(0, num_subseqs * num_steps, num_steps))
    # 在随机抽样的迭代过程中，
    # 来自两个相邻的、随机的、小批量中的子序列不一定在原始序列上相邻
    # 下面打乱子序列的起始索引
    random.shuffle(initial_indices)

    def data(pos):
        # 返回从pos位置开始的长度为num_steps的序列
        return corpus[pos: pos + num_steps]

    # num_batches 表示当前语料库可以产生几个批次。
    num_batches = num_subseqs // batch_size
    for i in range(0, batch_size * num_batches, batch_size):
        # 在这里，initial_indices包含子序列的随机起始索引
        # initial_indices_per_batch 则是当前批次的起始索引。
        initial_indices_per_batch = initial_indices[i: i + batch_size]
        # Y[j] 是 X[j] 后移一位的目标序列（预测下一个时间步）。
        X = [data(j) for j in initial_indices_per_batch]
        Y = [data(j + 1) for j in initial_indices_per_batch]
        yield torch.tensor(X), torch.tensor(Y)

class SeqDataLoader:  
    def __init__(self, txtpath,batch_size, num_steps, max_tokens=-1):
        """加载序列数据的迭代器
        Args:
            txtpath (str): 语料的txt文件路径
            batch_size (int): 批量大小，即每个批次的样本个数。
            num_steps (int): 每个样本(即子序列)的序列长度。
            max_tokens (int, optional): 提取语料库的最大长度。. Defaults to -1(不限制长度)
        """
        self.data_iter_fn = seq_data_iter_random
        self.corpus, self.vocab = load_corpus(txtpath,max_tokens)
        self.batch_size, self.num_steps = batch_size, num_steps

    '''
    Python中，迭代器是一种可以记住遍历位置的对象。
    __iter__(): 返回一个迭代器对象。
    __next__(): 返回容器中的下一个项目。

    迭代器与for语句：
    当我们写for x in fruits:时，实质上是for循环在背后调用iter(fruits)生成迭代器，再不断用next()取值，因此你可以遍历列表、字符串或其他可迭代对象。
    next()函数一项项地遍历数据，而无需一次性加载所有数据。

    如果 __iter__ 方法使用 yield 语句， yield 会使 __iter__ 成为一个生成器函数，调用时自动返回一个生成器对象，而生成器对象已经隐式实现了迭代器协议（即 __iter__ 和 __next__ 方法）。
    '''

    def __iter__(self):
        return self.data_iter_fn(self.corpus, self.batch_size, self.num_steps)

def load_data_iter(txtpath,batch_size, num_steps, max_tokens=-1):
    """返回数据集的迭代器和词表
    Args:
        txtpath (str): 语料的txt文件路径
        batch_size (int): 批量大小，即每个批次的样本个数。
        num_steps (int): 每个样本(即子序列)的序列长度。
        max_tokens (int, optional): 提取语料库的最大长度。. Defaults to -1(不限制长度)
    Returns:
        (tuple): (data_iter, vocab) data_iter数据集迭代器，每次返回一批num_steps长度的输出输出子序列，一批数量为batch_size。vocab 是词元字典，键是词元字符串，值是词元索引。
    """
    data_iter = SeqDataLoader(txtpath,
        batch_size, num_steps, max_tokens)
    return data_iter, data_iter.vocab

############    下面定义循环神经网络     ################




class RNNModelScratch: 
    """
    从零开始实现的循环神经网络模型.
    在技术场景中，“Scratch” 常用来描述 “不依赖高级封装，完全通过基础代码实现功能” .
    “Scratch” 还有其他含义：
    动词：
        “刮擦、划掉”：如 “scratch the surface”（只触及表面，未深入）。
        “临时取消”：如 “The meeting was scratched.”（会议被临时取消了）。
    名词：
        “刮痕、划痕”：如 a scratch on the screen（屏幕上的划痕）。
        “临时拼凑的东西”：如 a scratch team（临时组建的团队）。
    """

    def __init__(self, 
                 vocab_size, 
                 num_hiddens, 
                 device,
                 ):
        """初始化模型参数
        Args:
            vocab_size (int): 词表大小(即输入序列的维度)
            num_hiddens (int): 隐藏层大小
            device (_type_): 参数存储的位置。
            init_state (_type_): 初始状态
        """
        self.vocab_size, self.num_hiddens = vocab_size, num_hiddens
        self.params = self.get_params(vocab_size, num_hiddens, device)
        # 下面设置前向传播函数。
        self.forward_fn =  self.rnn

    def get_params(self, vocab_size, num_hiddens, device):
        """ 初始化模型参数（权重和偏置）。设输入序列xt，输出序列yt。一般输出序列与输入序列的元素形状是一样的。H(t)是隐状态。 H(t) = tanh( Xt*W_xh + H(t-1)*W_hh + b_h )，Y(t) = H(t)*W_hq + b_q
        Args:
            vocab_size (int): 词表大小(即输入序列的维度)
            num_hiddens (int): 隐藏层大小
            device (_type_): 参数存储的位置。
        Returns:
            list: 参数列表[W_xh, W_hh, b_h, W_hq, b_q]
        """
        num_inputs = num_outputs = vocab_size

        # 生成形状是shape的张量，张量元素初始化为均值0，标准差0.01的随机变量。变量放到device设备里。
        def normal(shape): 
            return torch.randn(size=shape, device=device) * 0.01

        # 隐藏层参数
        W_xh = normal((num_inputs, num_hiddens))
        W_hh = normal((num_hiddens, num_hiddens))
        b_h = torch.zeros(num_hiddens, device=device)
        # 输出层参数
        W_hq = normal((num_hiddens, num_outputs))
        b_q = torch.zeros(num_outputs, device=device)
        # 附加梯度
        params = [W_xh, W_hh, b_h, W_hq, b_q]
        for param in params:
            param.requires_grad_(True)
        return params

    def init_rnn_state(self, batch_size, num_hiddens, device):
        '''为 RNN 的第一个时间步创建初始隐藏状态（全零张量）。为了方便后续添加对象，用元组的形式返回。
        '''
        return (torch.zeros((batch_size, num_hiddens), device=device), )

    def rnn(self, inputs, state, params):
        """RNN 前向传播函数
        Args:
            inputs (tensor): 输入序列。形状(时间步数量，批量大小，词表大小)。要注意这里输入序列的尺寸。批量大小放到形状的第二个位置，可以方便矩阵计算：矩阵计算行之间独立，而不同批次也是独立的。
            state (tuple): 元组的元素是张量。表示初始隐藏状态，张量形状(批量大小，隐状态维度)。
            params (list): 参数列表[W_xh, W_hh, b_h, W_hq, b_q]
        Returns:
            tuple: ( 输出序列张量 , (state,) )  输出序列张量形状与inputs一致，state是该批次各个序列的当前状态。
        """

        # inputs的形状：(时间步数量，批量大小，词表大小)
        W_xh, W_hh, b_h, W_hq, b_q = params
        H, = state
        outputs = []
        # X的形状：(批量大小，词表大小)
        for X in inputs:
            H = torch.tanh(torch.mm(X, W_xh) + torch.mm(H, W_hh) + b_h)
            Y = torch.mm(H, W_hq) + b_q
            outputs.append(Y)
        return torch.cat(outputs, dim=0), (H,)

    '''
     Python 中，__call__ 是一个特殊方法（也称为 “魔术方法”），它允许一个类的实例像函数一样被调用。
    model = RNNModelScratch(...)  # 创建实例
    output, new_state = model(X, state)  # 直接调用实例，等价于 model.__call__(X, state)
    '''
    def __call__(self, X, state):
        X = F.one_hot(X.T, self.vocab_size).type(torch.float32)
        '''
        这段代码的作用是把输入张量X转换为独热编码.
        X.T会对输入张量X进行转置处理。假设X的形状是(batch_size, sequence_length)，转置之后就变成了(sequence_length, batch_size)。
        这么做的目的或许是为了让后续处理按照序列的时间步来进行。
        F.one_hot(X.T, self.vocab_size)会把转置后的张量X.T转换成独热编码形式。
        这里的self.vocab_size代表词汇表的大小，也就是独热编码向量的维度。
        转换完成后，张量的形状会变为(sequence_length, batch_size, self.vocab_size)。

        独热编码默认的数据类型是torch.long（也就是整数类型），但在深度学习模型里，通常需要使用浮点数类型。所以要把数据类型转换为torch.float32，这样才能作为模型的输入。
        '''
        
        
        return self.forward_fn(X, state, self.params)

    def begin_state(self, batch_size, device):
        return self.init_rnn_state(batch_size, self.num_hiddens, device)


class RNNModel(nn.Module):
    """循环神经网络模型，库函数实现"""
    def __init__(self, rnn_layer, vocab_size, **kwargs):
        super(RNNModel, self).__init__(**kwargs)
        self.rnn = rnn_layer
        self.vocab_size = vocab_size
        self.num_hiddens = self.rnn.hidden_size
        # 如果RNN是双向的（之后将介绍），num_directions应该是2，否则应该是1
        if not self.rnn.bidirectional:
            self.num_directions = 1
            self.linear = nn.Linear(self.num_hiddens, self.vocab_size)
        else:
            self.num_directions = 2
            self.linear = nn.Linear(self.num_hiddens * 2, self.vocab_size)

    def forward(self, inputs, state):
        X = F.one_hot(inputs.T.long(), self.vocab_size)
        X = X.to(torch.float32)
        Y, state = self.rnn(X, state)
        # 全连接层首先将Y的形状改为(时间步数*批量大小,隐藏单元数)
        # 它的输出形状是(时间步数*批量大小,词表大小)。
        output = self.linear(Y.reshape((-1, Y.shape[-1])))
        return output, state

    def begin_state(self, device, batch_size=1):
        if not isinstance(self.rnn, nn.LSTM):
            # nn.GRU以张量作为隐状态
            return  torch.zeros((self.num_directions * self.rnn.num_layers,
                                 batch_size, self.num_hiddens),
                                device=device)
        else:
            # nn.LSTM以元组作为隐状态
            return (torch.zeros((
                self.num_directions * self.rnn.num_layers,
                batch_size, self.num_hiddens), device=device),
                    torch.zeros((
                        self.num_directions * self.rnn.num_layers,
                        batch_size, self.num_hiddens), device=device))



if __name__ == '__main__':
    txtfile = "./Broken Butterflies.txt"

    if 0: # 测试语料库预处理
        lines = read_txtfile(txtfile)

        print(f'# 文本总行数: {len(lines)}')
        print(lines[0])
        # print(lines[1])

        tokens = tokenize(lines)
        conter = count_corpus(tokens)
        print(' counts of the : ',conter['the'])

        voc = Vocab()
        #字典(Dictionary) items() 函数以列表返回可遍历的(键, 值) 元组数组。
        print(list(voc.token_to_idx.items())[:10]) 
        print('文本:', tokens[0])
        print('索引:', voc[tokens[0]])
        print(len(voc))

        voc.set_txtfile(txtfile)

        print(list(voc.token_to_idx.items())[:10])
        print('文本:', tokens[0])
        print('索引:', voc[tokens[0]])
        print(len(voc))

        my_seq = list(range(35))
        for X, Y in seq_data_iter_random(my_seq, batch_size=2, num_steps=5):
            print('X: ', X, '\nY:', Y)

    if 0: # 测试自定义cnn

        print('尝试打印2个10维的热独编码， 分别在4，5的位置是1')
        print(F.one_hot(torch.tensor([4, 5]), 10))

        batch_size, num_steps = 32, 35
        train_iter, vocab = load_data_iter(txtfile,batch_size, num_steps)
        
        for batch_input, batch_output in train_iter:
            
            # 迭代器返回的是 (批量大小，时间步数量，词表大小)
            X = batch_input
            # X = batch_input.permute(1, 0)  # 将原来形状[0,1]的张量变成[1,0]的张量。 RNNModelScratch 自定义的 __call__函数已经做了热独编码的处理了。

            num_hiddens = 512
            net = RNNModelScratch(len(vocab), num_hiddens, TrainAndTest.try_gpu(),)
            state = net.begin_state(batch_size, TrainAndTest.try_gpu())
            Y, new_state = net(X.to(TrainAndTest.try_gpu()), state)
            print('检查输出是否具有正确的形状')
            print(Y.shape)
            print(new_state[0].shape)
            print(X.shape)
            print(len(vocab))
            break

